"""
视频处理模块
处理RTSP视频流和帧率控制
"""

import cv2
import numpy as np
import threading
import time
import logging
from typing import Optional, Callable, Tuple
from queue import Queue, Empty
import os

class RTSPProcessor:
    """RTSP视频流处理器"""
    
    def __init__(self, rtsp_url: str, fps_limit: int = 15, 
                 buffer_size: int = 5, resize_shape: Optional[Tuple[int, int]] = None):
        """
        初始化RTSP处理器
        
        Args:
            rtsp_url: RTSP流地址
            fps_limit: 帧率限制
            buffer_size: 缓冲区大小
            resize_shape: 调整尺寸 (width, height)
        """
        self.rtsp_url = rtsp_url
        self.fps_limit = fps_limit
        self.buffer_size = buffer_size
        self.resize_shape = resize_shape
        
        # 视频捕获对象
        self.cap = None
        self.frame_queue = Queue(maxsize=buffer_size)
        
        # 控制标志
        self.is_running = False
        self.capture_thread = None
        
        # 统计信息
        self.frame_count = 0
        self.fps_actual = 0.0
        self.last_fps_time = time.time()
        self.last_frame_time = 0.0
        
        # 重连设置
        self.max_reconnect_attempts = 5
        self.reconnect_delay = 2.0
        
        # 日志
        self.logger = logging.getLogger(__name__)
        
    def connect(self) -> bool:
        """连接到RTSP流"""
        try:
            self.cap = cv2.VideoCapture(self.rtsp_url)
            
            # 设置缓冲区大小
            self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            
            # 测试连接
            ret, frame = self.cap.read()
            if not ret or frame is None:
                self.logger.error("无法从RTSP流读取帧")
                return False
            
            # 获取视频属性
            width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            fps = self.cap.get(cv2.CAP_PROP_FPS)
            
            self.logger.info(f"RTSP连接成功: {width}x{height} @{fps:.1f}fps")
            return True
            
        except Exception as e:
            self.logger.error(f"RTSP连接失败: {e}")
            return False
    
    def disconnect(self):
        """断开连接"""
        if self.cap:
            self.cap.release()
            self.cap = None
        self.logger.info("RTSP连接已断开")
    
    def start_capture(self):
        """开始捕获视频流"""
        if not self.connect():
            return False
        
        self.is_running = True
        self.capture_thread = threading.Thread(target=self._capture_loop, daemon=True)
        self.capture_thread.start()
        
        self.logger.info("视频捕获已开始")
        return True
    
    def stop_capture(self):
        """停止捕获视频流"""
        self.is_running = False
        
        if self.capture_thread and self.capture_thread.is_alive():
            self.capture_thread.join(timeout=2.0)
        
        self.disconnect()
        
        # 清空队列
        while not self.frame_queue.empty():
            try:
                self.frame_queue.get_nowait()
            except Empty:
                break
        
        self.logger.info("视频捕获已停止")
    
    def _capture_loop(self):
        """视频捕获循环"""
        frame_interval = 1.0 / self.fps_limit
        reconnect_count = 0
        
        while self.is_running:
            try:
                if not self.cap or not self.cap.isOpened():
                    # 尝试重连
                    if reconnect_count < self.max_reconnect_attempts:
                        self.logger.warning(f"尝试重连 ({reconnect_count + 1}/{self.max_reconnect_attempts})")
                        time.sleep(self.reconnect_delay)
                        
                        if self.connect():
                            reconnect_count = 0
                        else:
                            reconnect_count += 1
                            continue
                    else:
                        self.logger.error("达到最大重连次数，停止捕获")
                        break
                
                # 读取帧
                ret, frame = self.cap.read()
                if not ret or frame is None:
                    self.logger.warning("读取帧失败")
                    reconnect_count += 1
                    continue
                
                # 重置重连计数
                reconnect_count = 0
                
                # 调整帧大小
                if self.resize_shape:
                    frame = cv2.resize(frame, self.resize_shape)
                
                # 控制帧率
                current_time = time.time()
                if current_time - self.last_frame_time < frame_interval:
                    continue
                
                self.last_frame_time = current_time
                
                # 添加到队列
                if self.frame_queue.full():
                    try:
                        self.frame_queue.get_nowait()  # 移除旧帧
                    except Empty:
                        pass
                
                self.frame_queue.put(frame)
                
                # 更新统计信息
                self._update_fps_stats()
                
            except Exception as e:
                self.logger.error(f"捕获循环错误: {e}")
                time.sleep(0.1)
    
    def _update_fps_stats(self):
        """更新FPS统计"""
        self.frame_count += 1
        current_time = time.time()
        
        if current_time - self.last_fps_time >= 1.0:
            self.fps_actual = self.frame_count / (current_time - self.last_fps_time)
            self.frame_count = 0
            self.last_fps_time = current_time
    
    def get_frame(self, timeout: float = 1.0) -> Optional[np.ndarray]:
        """
        获取最新帧
        
        Args:
            timeout: 超时时间
            
        Returns:
            视频帧或None
        """
        try:
            return self.frame_queue.get(timeout=timeout)
        except Empty:
            return None
    
    def get_stats(self) -> dict:
        """获取统计信息"""
        return {
            "is_running": self.is_running,
            "fps_actual": self.fps_actual,
            "fps_limit": self.fps_limit,
            "queue_size": self.frame_queue.qsize(),
            "connected": self.cap is not None and self.cap.isOpened() if self.cap else False
        }


class VideoSaver:
    """视频保存器"""
    
    def __init__(self, output_path: str, fps: float = 15.0, 
                 frame_size: Tuple[int, int] = (1280, 720)):
        """
        初始化视频保存器
        
        Args:
            output_path: 输出文件路径
            fps: 保存帧率
            frame_size: 帧尺寸
        """
        self.output_path = output_path
        self.fps = fps
        self.frame_size = frame_size
        
        # 创建输出目录
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # 视频写入器
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        self.writer = cv2.VideoWriter(output_path, fourcc, fps, frame_size)
        
        self.frame_count = 0
        self.is_recording = False
        
        # 日志
        self.logger = logging.getLogger(__name__)
    
    def start_recording(self):
        """开始录制"""
        self.is_recording = True
        self.logger.info(f"开始录制视频: {self.output_path}")
    
    def stop_recording(self):
        """停止录制"""
        self.is_recording = False
        if self.writer:
            self.writer.release()
        self.logger.info(f"录制完成，共{self.frame_count}帧")
    
    def write_frame(self, frame: np.ndarray):
        """写入帧"""
        if not self.is_recording or not self.writer:
            return
        
        # 调整帧大小
        if frame.shape[:2][::-1] != self.frame_size:
            frame = cv2.resize(frame, self.frame_size)
        
        self.writer.write(frame)
        self.frame_count += 1
    
    def __del__(self):
        """析构函数"""
        if hasattr(self, 'writer') and self.writer:
            self.writer.release()


class FrameProcessor:
    """帧处理器基类"""
    
    def process_frame(self, frame: np.ndarray) -> np.ndarray:
        """
        处理帧
        
        Args:
            frame: 输入帧
            
        Returns:
            处理后的帧
        """
        raise NotImplementedError
    
    def get_stats(self) -> dict:
        """获取处理统计信息"""
        return {}


class FPSCounter:
    """FPS计数器"""
    
    def __init__(self, window_size: int = 30):
        """
        初始化FPS计数器
        
        Args:
            window_size: 滑动窗口大小
        """
        self.window_size = window_size
        self.frame_times = []
        self.last_time = time.time()
    
    def update(self) -> float:
        """更新并返回当前FPS"""
        current_time = time.time()
        self.frame_times.append(current_time - self.last_time)
        self.last_time = current_time
        
        # 保持窗口大小
        if len(self.frame_times) > self.window_size:
            self.frame_times.pop(0)
        
        # 计算平均FPS
        if len(self.frame_times) > 0:
            avg_frame_time = sum(self.frame_times) / len(self.frame_times)
            return 1.0 / avg_frame_time if avg_frame_time > 0 else 0.0
        
        return 0.0
    
    def get_fps(self) -> float:
        """获取当前FPS"""
        if len(self.frame_times) > 0:
            avg_frame_time = sum(self.frame_times) / len(self.frame_times)
            return 1.0 / avg_frame_time if avg_frame_time > 0 else 0.0
        return 0.0


if __name__ == "__main__":
    # 测试代码
    logging.basicConfig(level=logging.INFO)
    
    # 测试RTSP处理器
    rtsp_url = "rtsp://admin:Lorytech2015@192.168.3.79:554/cam/realmonitor?channel=1&subtype=0"
    processor = RTSPProcessor(rtsp_url, fps_limit=15, resize_shape=(1280, 720))
    
    # 开始捕获
    if processor.start_capture():
        fps_counter = FPSCounter()
        
        print("按 'q' 键退出...")
        
        while True:
            frame = processor.get_frame(timeout=1.0)
            if frame is None:
                print("未获取到帧")
                continue
            
            # 更新FPS
            fps = fps_counter.update()
            
            # 显示统计信息
            stats = processor.get_stats()
            cv2.putText(frame, f"FPS: {fps:.1f}", (10, 30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            cv2.putText(frame, f"Queue: {stats['queue_size']}", (10, 60), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            
            cv2.imshow("RTSP Test", frame)
            
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        
        processor.stop_capture()
        cv2.destroyAllWindows()
    else:
        print("无法连接到RTSP流") 